/*
Cuckoo Sandbox - Automated Malware Analysis.
Copyright (C) 2016-2017 Cuckoo Foundation.

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include <stdio.h>
#include <stdint.h>
#include <windows.h>
#include "hooking.h"
#include "memory.h"
#include "misc.h"
#include "log.h"
#include "native.h"
#include "pipe.h"

static array_t g_last_guard_page_addr_array;
static slab_t g_guard_page_referer_slab;
static uintptr_t g_guard_pages_wl[32];
static uint32_t g_guard_page_wl_count;

static dnq_t g_guard_pages;

void find_guard_page_whitelist(uintptr_t *addresses, uint32_t *count);

void exploit_init()
{
    LDR_MODULE *mod; PEB *peb = get_peb();

    array_init(&g_last_guard_page_addr_array);
    slab_init(
        &g_guard_page_referer_slab, sizeof(hook_t), 128,
        PAGE_EXECUTE_READWRITE
    );

#if !EXPLOIT_GUARD_SUPPORT_ENABLED
    return;
#endif

    find_guard_page_whitelist(g_guard_pages_wl, &g_guard_page_wl_count);

    static void *guard_pages[0x1000]; uint32_t guard_page_count = 0;

    for (mod = (LDR_MODULE *) peb->LoaderData->InLoadOrderModuleList.Flink;
         mod->BaseAddress != NULL;
         mod = (LDR_MODULE *) mod->InLoadOrderModuleList.Flink) {

        guard_pages[guard_page_count++] = mod->BaseAddress;
    }

    dnq_init(
        &g_guard_pages, guard_pages, sizeof(uintptr_t), guard_page_count
    );

    for (uint32_t idx = 0; idx < guard_page_count; idx++) {
        exploit_set_guard_page((void *) dnq_iterptr(&g_guard_pages)[idx]);
    }
}

int exploit_is_registered_guard_page(uintptr_t addr)
{
    return dnq_hasptr(&g_guard_pages, addr & ~0xfff);
}

int WINAPI exploit_set_guard_page(void *addr)
{
    if(exploit_is_registered_guard_page((uintptr_t) addr) == 0) {
        return -1;
    }

    MEMORY_BASIC_INFORMATION_CROSS mbi;

    addr = (void *)((uintptr_t) addr & ~0xfff);
    if(virtual_query(addr, &mbi) == FALSE ||
            (mbi.Protect & PAGE_READABLE) == 0) {
        return -1;
    }

    if(NT_SUCCESS(virtual_protect(
            addr, 0x1000, mbi.Protect | PAGE_GUARD)) == FALSE) {
        return -1;
    }
    return 0;
}

int WINAPI exploit_unset_guard_page(void *addr)
{
    if(exploit_is_registered_guard_page((uintptr_t) addr) == 0) {
        return -1;
    }

    MEMORY_BASIC_INFORMATION_CROSS mbi;

    addr = (void *)((uintptr_t) addr & ~0xfff);
    if(virtual_query(addr, &mbi) == FALSE ||
            (mbi.Protect & PAGE_READABLE) == 0) {
        return -1;
    }

    if(NT_SUCCESS(virtual_protect(
            addr, 0x1000, mbi.Protect & ~PAGE_GUARD)) == FALSE) {
        return -1;
    }
    return 0;
}

void *exploit_get_last_guard_page()
{
    void *ret = array_get(
        &g_last_guard_page_addr_array, get_current_thread_id()
    );
    array_set(
        &g_last_guard_page_addr_array, get_current_thread_id(), NULL
    );
    return ret;
}

void exploit_set_last_guard_page(void *addr)
{
    array_set(
        &g_last_guard_page_addr_array, get_current_thread_id(), addr
    );
}

static int addr_in_stack(void *addr)
{
    NT_TIB *tib = (NT_TIB *) readtls(TLS_TEB);
    return addr >= tib->StackLimit && addr < tib->StackBase;
}

// Currently we just check if the memory block is of MEM_PRIVATE type.
// There are some APIs for traversing the heap, but they might miss
// some of the heap blocks (e.g., regions allocated through VirtualAlloc).
static int addr_in_heap(void *addr)
{
    MEMORY_BASIC_INFORMATION_CROSS mbi;
    virtual_query(addr, &mbi);
    return mbi.Type == MEM_PRIVATE;
}

int exploit_is_stack_pivoted()
{
    void *ptr = NULL;
    return addr_in_stack(&ptr) == 0;
}

int exploit_makes_stack_executable(
    HANDLE process_handle, PVOID addr, DWORD new_protection)
{
    return (new_protection & PAGE_EXECUTABLE) != 0 &&
        addr_in_stack(addr) != FALSE &&
        process_handle == get_current_process();
}

int exploit_makes_heap_executable(
    HANDLE process_handle, PVOID addr, DWORD new_protection)
{
    return (new_protection & PAGE_EXECUTABLE) != 0 &&
        addr_in_heap(addr) != FALSE &&
        process_handle == get_current_process();
}

static funcoff_t _LdrGetProcedureAddressEx[] = {
    {0x4ce7ba58, 0x302b8, 0},
    {0, 0, 0},
};

static funcoff_t _LdrFindResource_U[] = {
    {0x4ce7ba58, 0x41f53, 0},
    {0, 0, 0},
};

static funcoff_t _LdrLoadDll[] = {
    {0x4ce7ba58, 0x3c4b5, 0},
    {0, 0, 0},
};

static funcoff_t _RtlImageNtHeader[] = {
    {0x4ce7ba58, 0x3317f, 0},
    {0, 0, 0},
};

static funcoff_t _RtlPcToFileHeader[] = {
    {0x4ce7ba58, 0x400ff, 0},
    {0, 0, 0},
};

static mod2funcoff_t _ntdll[] = {
    {"LdrGetProcedureAddressEx", _LdrGetProcedureAddressEx},
    {"LdrFindResource_U", _LdrFindResource_U},
    {"LdrLoadDll", _LdrLoadDll},
    {"RtlImageNtHeader", _RtlImageNtHeader},
    {"RtlPcToFileHeader", _RtlPcToFileHeader},
    {NULL, NULL},
};

void find_guard_page_whitelist(uintptr_t *addresses, uint32_t *count)
{
    uint8_t *modaddr = (uint8_t *) GetModuleHandle("ntdll");
    for (mod2funcoff_t *p = _ntdll; p->funcname != NULL; p++) {
        uint8_t *ret = module_addr_timestamp(modaddr, 0, p->funcoff, NULL);
        if(ret != NULL) {
            addresses[*count] = (uintptr_t) ret;
            *count += 1;
        }
    }

    if(*count == 0) {
        log_action("gatherer");
    }
}

int exploit_is_guard_page_referer_whitelisted(
    uintptr_t *addrs, uint32_t count)
{
    // If we were unable to determine the Guard Page whitelist this means that
    // we can't use the feature on this OS yet, as it would generate a lot of
    // false positives (for legitimate calls to GetProcAddress etc).
    if(g_guard_page_wl_count == 0) {
        return 1;
    }

    for (uint32_t idx = 0; idx < count; idx++) {
        for (uint32_t jdx = 0; jdx < g_guard_page_wl_count; jdx++) {
            if(g_guard_pages_wl[jdx] == addrs[idx]) {
                return 1;
            }
        }
    }
    return 0;
}

int exploit_insn_rewrite_to_lea(uint8_t *buf, uint8_t *insn)
{
    if(*insn == 0x66) {
        insn++;
    }

    char hex[40];

    if(*insn == 0x0f) {
        switch (*++insn) {
        case 0xb6: case 0xb7:
            insn++;
            break;

        default:
            hexdump(hex, insn-1, 16);
            pipe("INFO:unhandled insn %z", hex);
            return -1;
        }
    }
    else {
        switch (*insn) {
        case 0x39: case 0x3b: case 0x80: case 0x83: case 0x85: case 0x8a:
        case 0x8b:
            insn++;
            break;

        case 0xf3:
            // "rep movsd" turns into the guard page "esi"
            if(*++insn == 0xa5) {
                *buf++ = 0x8d;
                *buf++ = 0x06;
                return 2;
            }

            // Intentional fall-through.

        default:
            hexdump(hex, insn, 16);
            pipe("INFO:unhandled insn %z", hex);
            return -1;
        }
    }

    uint8_t modrm, rm, sib, reg, mult = 0; uint32_t imm = 0;

    modrm = *insn++;
    rm = modrm & 7;

    if((modrm & 0xc0) == 0xc0) {
        return -1;
    }

    if((modrm & 0x07) == 0x04) {
        sib = *insn++;
        rm = sib & 7;
        reg = (sib >> 3) & 7;
        mult = 1 + (sib >> 6);
    }

    if((modrm & 0xc0) == 0x40) {
        imm = *(int8_t *) insn;
    }

    if((modrm & 0xc0) == 0x80) {
        imm = *(int32_t *) insn;
    }

    *buf++ = 0x8d;
    if(mult == 0) {
        *buf++ = 0x80 + rm;
    }
    else {
        *buf++ = 0x84;
        *buf++ = ((mult - 1) << 6) + (reg << 3) + rm;
    }

    *(uint32_t *) buf = imm;
    return 6 + (mult != 0);
}

int exploit_hotpatch_guard_page_referer(uintptr_t pc)
{
    uint8_t *insn = (uint8_t *) pc;

    hook_t *h = slab_getmem(&g_guard_page_referer_slab);

    h->type = HOOK_TYPE_GUARD;
    h->addr = insn;
    h->library = "__guard__";

    int r = hook(h, NULL);
    if(r < 0) {
        slab_return_last(&g_guard_page_referer_slab);
        return r;
    }
    return h->stub_used;
}
